package com.ruoyi.common.utils.ps2;
import org.apache.poi.ss.usermodel.*;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.logging.*;

public class OptimizedImageDownloader{
    // 配置参数（可根据需要调整）
    private static final class Config {
        static final String BASE_URL = "http://prc-videoframe-pub-1258344703.cos.ap-guangzhou.myqcloud.com/ad_creative_engine/text2image_pipeline_prd/";
        static final String INPUT_EXCEL = "C:\\Users\\xinhuiRong\\Desktop\\新建文件夹\\人物属性.xlsx";
        static final String OUTPUT_DIR = "H:\\person";
        static final String PROGRESS_FILE = "H:\\person\\download_progress.txt";
        static final String FAILED_RECORDS_FILE = "H:\\person\\failed_records.xlsx";
        static final String LOG_DIR = "H:\\person\\logs";
        static final String FAILED_LOG = "H:\\person\\logs\\failed_downloads.log";

        static final int BATCH_SIZE = 1000;
        static final int THREAD_POOL_SIZE = 8;
        static final int CONNECT_TIMEOUT = 15000;
        static final int READ_TIMEOUT = 45000;
        static final int MAX_RETRY = 5;
        static final int MEMORY_CHECK_INTERVAL = 50;
    }

    private static final Logger logger = Logger.getLogger(OptimizedImageDownloader.class.getName());

    // 数据存储
    private static final List<Map<String, String>> allRecords = new ArrayList<>();
    private static final Set<Integer> processedIndices = ConcurrentHashMap.newKeySet();
    private static final Set<String> downloadedFiles = ConcurrentHashMap.newKeySet();
    private static final Set<Integer> failedIndices = ConcurrentHashMap.newKeySet();

    public static void main(String[] args) {
        setupLogger();

        try {
            // 加载Excel数据
            loadExcelData();
            logger.info("Loaded " + allRecords.size() + " records from Excel");

            // 加载进度和已下载文件
            loadProgress();
            scanDownloadedFiles();

            // 下载图片
            downloadImages();

            logger.info("All images downloaded successfully!");
        } catch (Exception e) {
            logger.log(Level.SEVERE, "Fatal error: " + e.getMessage(), e);
        } finally {
            saveFailedRecords();
            cleanup();
        }
    }

    private static void setupLogger() {
        try {
            Files.createDirectories(Paths.get(Config.LOG_DIR));

            FileHandler fileHandler = new FileHandler(Config.LOG_DIR + "/downloader.log", true);
            fileHandler.setFormatter(new SimpleFormatter() {
                private static final String format = "[%1$tF %1$tT] [%2$-7s] %3$s %n";

                @Override
                public synchronized String format(LogRecord lr) {
                    return String.format(format,
                            new Date(lr.getMillis()),
                            lr.getLevel().getLocalizedName(),
                            lr.getMessage());
                }
            });
            logger.addHandler(fileHandler);

            ConsoleHandler consoleHandler = new ConsoleHandler();
            consoleHandler.setFormatter(new SimpleFormatter());
            logger.addHandler(consoleHandler);

            logger.setLevel(Level.ALL);
        } catch (IOException e) {
            System.err.println("Failed to setup logger: " + e.getMessage());
        }
    }

    private static void loadExcelData() throws IOException {
        try (InputStream is = new FileInputStream(Config.INPUT_EXCEL);
             Workbook workbook = new XSSFWorkbook(is)) {

            Sheet sheet = workbook.getSheetAt(0);
            Row headerRow = sheet.getRow(0);

            // 读取表头
            List<String> headers = new ArrayList<>();
            for (Cell cell : headerRow) {
                headers.add(cell.getStringCellValue());
            }

            // 读取数据行
            for (int i = 1; i <= sheet.getLastRowNum(); i++) {
                Row row = sheet.getRow(i);
                if (row == null) continue;

                Map<String, String> record = new LinkedHashMap<>();
                for (int j = 0; j < headers.size(); j++) {
                    Cell cell = row.getCell(j, Row.MissingCellPolicy.CREATE_NULL_AS_BLANK);
                    String value = "";

                    switch (cell.getCellType()) {
                        case STRING:
                            value = cell.getStringCellValue();
                            break;
                        case NUMERIC:
                            value = String.valueOf((int) cell.getNumericCellValue());
                            break;
                        case BOOLEAN:
                            value = String.valueOf(cell.getBooleanCellValue());
                            break;
                        default:
                            value = "";
                    }

                    record.put(headers.get(j), value.trim());
                }

                if (!record.get("md5").isEmpty()) {
                    allRecords.add(record);
                }
            }
        }
    }

    private static void loadProgress() {
        try {
            if (Files.exists(Paths.get(Config.PROGRESS_FILE))) {
                Files.lines(Paths.get(Config.PROGRESS_FILE))
                        .filter(line -> !line.trim().isEmpty())
                        .mapToInt(Integer::parseInt)
                        .forEach(processedIndices::add);
                logger.info("Loaded progress: " + processedIndices.size() + " processed indices");
            }
        } catch (Exception e) {
            logger.log(Level.WARNING, "Failed to load progress: " + e.getMessage());
        }
    }

    private static void scanDownloadedFiles() {
        try {
            Path outputDir = Paths.get(Config.OUTPUT_DIR);
            if (Files.exists(outputDir)) {
                Files.walk(outputDir)
                        .filter(Files::isRegularFile)
                        .map(path -> path.getFileName().toString())
                        .forEach(downloadedFiles::add);
                logger.info("Found " + downloadedFiles.size() + " existing files");
            }
        } catch (IOException e) {
            logger.log(Level.WARNING, "Failed to scan downloaded files: " + e.getMessage());
        }
    }

    private static void downloadImages() {
        try {
            Files.createDirectories(Paths.get(Config.OUTPUT_DIR));
        } catch (IOException e) {
            logger.log(Level.SEVERE, "Failed to create output directory: " + e.getMessage());
            return;
        }

        ExecutorService executor = Executors.newFixedThreadPool(Config.THREAD_POOL_SIZE);
        CompletionService<Void> completionService = new ExecutorCompletionService<>(executor);

        int totalTasks = 0;
        int completedTasks = 0;

        // 提交下载任务
        for (int i = 0; i < allRecords.size(); i++) {
            Map<String, String> record = allRecords.get(i);
            String md5 = record.get("md5");

            if (processedIndices.contains(i) || md5.isEmpty()) {
                continue;
            }

            int finalI = i;
            completionService.submit(() -> {
                try {
                    downloadImageWithRetry(finalI, record);
                    return null;
                } catch (Exception e) {
                    logger.log(Level.WARNING, "Task failed (index " + finalI + "): " + e.getMessage());
                    throw e;
                }
            });
            totalTasks++;
        }

        // 等待任务完成
        while (completedTasks < totalTasks) {
            try {
                Future<Void> future = completionService.take();
                future.get();
                completedTasks++;

                // 定期保存进度
                if (completedTasks % Config.MEMORY_CHECK_INTERVAL == 0) {
                    saveProgress();
                    System.gc();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            } catch (ExecutionException e) {
                logger.log(Level.WARNING, "Task execution error: " + e.getCause().getMessage());
            }
        }

        shutdownExecutor(executor);
        saveProgress();
    }

    private static void downloadImageWithRetry(int index, Map<String, String> record) {
        String md5 = record.get("md5");
        String fileName = md5 + ".jpeg";
        String imageUrl = Config.BASE_URL + fileName;

        // 跳过已下载文件
        if (downloadedFiles.contains(fileName)) {
            markAsProcessed(index);
            return;
        }

        // 计算批次目录
        int batchNumber = index / Config.BATCH_SIZE;
        int startRange = batchNumber * Config.BATCH_SIZE;
        int endRange = startRange + Config.BATCH_SIZE - 1;
        String batchDirName = startRange + "-" + endRange;

        try {
            // 创建批次目录
            Path batchDir = Paths.get(Config.OUTPUT_DIR, batchDirName);
            synchronized (OptimizedImageDownloader.class) {
                if (!Files.exists(batchDir)) {
                    Files.createDirectories(batchDir);
                }
            }

            Path outputPath = batchDir.resolve(fileName);

            // 重试下载
            for (int attempt = 1; attempt <= Config.MAX_RETRY; attempt++) {
                try {
                    if (downloadSingleImage(imageUrl, outputPath)) {
                        downloadedFiles.add(fileName);
                        markAsProcessed(index);
                        return;
                    }
                } catch (IOException e) {
                    logger.log(Level.WARNING, String.format("Attempt %d/%d failed for %s: %s",
                            attempt, Config.MAX_RETRY, imageUrl, e.getMessage()));
                }

                if (attempt < Config.MAX_RETRY) {
                    Thread.sleep(2000 * attempt); // 指数退避
                }
            }

            // 下载失败处理
            logger.severe("Download failed after " + Config.MAX_RETRY + " attempts: " + imageUrl);
            failedIndices.add(index);

            // 记录失败日志
            Files.write(Paths.get(Config.FAILED_LOG),
                    (imageUrl + System.lineSeparator()).getBytes(),
                    StandardOpenOption.CREATE, StandardOpenOption.APPEND);

        } catch (Exception e) {
            logger.log(Level.SEVERE, "Error processing download (index " + index + "): " + e.getMessage());
            failedIndices.add(index);
        }
    }

    private static boolean downloadSingleImage(String imageUrl, Path outputPath) throws IOException {
        HttpURLConnection connection = null;
        try {
            URL url = new URL(imageUrl);
            connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("GET");
            connection.setConnectTimeout(Config.CONNECT_TIMEOUT);
            connection.setReadTimeout(Config.READ_TIMEOUT);
            connection.setInstanceFollowRedirects(true);

            int responseCode = connection.getResponseCode();
            if (responseCode == HttpURLConnection.HTTP_OK) {
                try (InputStream in = connection.getInputStream();
                     OutputStream out = new FileOutputStream(outputPath.toFile())) {
                    byte[] buffer = new byte[16384]; // 16KB buffer
                    int bytesRead;
                    while ((bytesRead = in.read(buffer)) != -1) {
                        out.write(buffer, 0, bytesRead);
                    }
                }
                return true;
            } else {
                logger.warning("HTTP " + responseCode + " for URL: " + imageUrl);
                return false;
            }
        } finally {
            if (connection != null) {
                connection.disconnect();
            }
        }
    }

    private static synchronized void markAsProcessed(int index) {
        processedIndices.add(index);
    }

    private static void saveProgress() {
        try {
            Path tempFile = Paths.get(Config.PROGRESS_FILE + ".tmp");
            try (BufferedWriter writer = Files.newBufferedWriter(tempFile)) {
                processedIndices.stream()
                        .sorted()
                        .map(String::valueOf)
                        .forEach(line -> {
                            try {
                                writer.write(line);
                                writer.newLine();
                            } catch (IOException e) {
                                logger.log(Level.WARNING, "Error writing progress: " + e.getMessage());
                            }
                        });
            }
            Files.move(tempFile, Paths.get(Config.PROGRESS_FILE), StandardCopyOption.REPLACE_EXISTING);
        } catch (IOException e) {
            logger.log(Level.WARNING, "Failed to save progress: " + e.getMessage());
        }
    }

    private static void saveFailedRecords() {
        if (failedIndices.isEmpty()) {
            return;
        }

        try (Workbook workbook = new XSSFWorkbook()) {
            Sheet sheet = workbook.createSheet("Failed Downloads");

            // 创建表头（原表头 + 状态列）
            Row headerRow = sheet.createRow(0);
            int colNum = 0;
            for (String header : allRecords.get(0).keySet()) {
                headerRow.createCell(colNum++).setCellValue(header);
            }
            headerRow.createCell(colNum).setCellValue("Download Status");

            // 填充失败记录
            int rowNum = 1;
            for (int index : failedIndices) {
                Map<String, String> record = allRecords.get(index);
                Row row = sheet.createRow(rowNum++);

                colNum = 0;
                for (String value : record.values()) {
                    row.createCell(colNum++).setCellValue(value);
                }
                row.createCell(colNum).setCellValue("Failed");
            }

            // 保存文件
            try (FileOutputStream out = new FileOutputStream(Config.FAILED_RECORDS_FILE)) {
                workbook.write(out);
            }

            logger.warning("Saved " + failedIndices.size() + " failed records to " + Config.FAILED_RECORDS_FILE);
        } catch (IOException e) {
            logger.log(Level.SEVERE, "Failed to save failed records: " + e.getMessage(), e);
        }
    }

    private static void shutdownExecutor(ExecutorService executor) {
        executor.shutdown();
        try {
            if (!executor.awaitTermination(60, TimeUnit.SECONDS)) {
                executor.shutdownNow();
                if (!executor.awaitTermination(60, TimeUnit.SECONDS)) {
                    logger.warning("Executor did not terminate");
                }
            }
        } catch (InterruptedException e) {
            executor.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    private static void cleanup() {
        try {
            if (processedIndices.size() == allRecords.size()) {
                Files.deleteIfExists(Paths.get(Config.PROGRESS_FILE));
                logger.info("Cleanup completed");
            }
        } catch (IOException e) {
            logger.log(Level.WARNING, "Cleanup failed: " + e.getMessage());
        }
    }
}